# -*- coding: UTF-8 -*-
'''
@Project ：push_rk
@File ：myHandle.py
@IDE ：PyCharm
@Author ：苦瓜
@Date ：2025/10/10 10:27
@Note: Something beautiful is about to happen !
'''
import matplotlib.pyplot as plt
#   	导入必要的库和模块
from sklearn.datasets import load_boston
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import torch
from torch import nn

#   	加载波士顿房价数据集，并对数据进行标准化处理
data, target = load_boston(return_X_y=True)
data = StandardScaler().fit_transform(data)
target = StandardScaler().fit_transform(target.reshape(-1, 1))

#  .	定义时间窗口大小，并将数据转换为适合LSTM模型输入的格式
c = 7
x_list = []
y_list = []
for i in range(len(data)-c):
    x_list.append(data[i:i+c])
    y_list.append(target[i+c])

# 类型转换
x_list = torch.tensor(x_list, dtype=torch.float32)
y_list = torch.tensor(y_list, dtype=torch.float32)

#  .	划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(x_list, y_list, shuffle=False, random_state=42)

# 设置输入形状
input_size = X_train.shape[2]

#  .	定义一个LSTM模型，包含两个LSTM层和一个全连接层
class MyLSTM(nn.Module):

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.rnn1 = nn.LSTM(input_size=input_size, hidden_size=10, batch_first=True)
        self.rnn2 = nn.LSTM(input_size=10, hidden_size=10, batch_first=True)
        self.fc = nn.Linear(in_features=10, out_features=1)

    def forward(self, inputs):
        #  .	前向传播函数，定义数据如何通过模型
        x, _ = self.rnn1(inputs)
        x, _ = self.rnn2(x)
        x = self.fc(x[:, -1, :])
        return x

if __name__ == '__main__':
    #  .	实例化LSTM模型
    model_lstm = MyLSTM()
    #  .	使用MSE损失函数和Adam优化器来训练模型
    cost = nn.MSELoss()
    optimizer = torch.optim.Adam(model_lstm.parameters(), lr=0.001)
    losses_ = []
    model_lstm.train()
    for epoch in range(1000):
        optimizer.zero_grad()
        y_pred_train = model_lstm(X_train)
        loss_ = cost(y_pred_train, y_train)
        losses_.append(loss_.item())
        loss_.backward()
        optimizer.step()
        if epoch % 10 == 0 :
            print(f"<Model ({epoch}) train loss : ({loss_})>")

    model_lstm.eval()
    #  .	对测试集进行预测
    with torch.no_grad():
        y_pred_test = model_lstm(X_test)
    print(f"""
        === 测试集预测结果 ===
        {y_pred_test}
    """)

    # 10.	将预测结果和真实值进行可视化对比
    plt.plot(y_pred_test, label="y pred test ", c='g')
    plt.plot(y_test, label="y test ", c='r')
    plt.legend()
    plt.show()

    # 损失图
    plt.plot(losses_)
    plt.title("loss plot")
    plt.show()
